#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
简化版数据访问工具
Simplified Data Access Tool

Author: ML Team
Date: 2025-11-16
Description: 便捷访问v1和v2版本的处理后数据
"""

import pandas as pd
import os
from typing import Tuple


class SimpleDataManager:
    """简化的数据管理器"""

    def __init__(self):
        """初始化数据管理器"""
        self.data_dir = os.path.dirname(os.path.abspath(__file__))

    def list_versions(self):
        """列出可用版本"""
        print("="*50)
        print("可用的数据版本")
        print("="*50)

        versions = []

        # 检查v1版本
        v1_train = os.path.join(self.data_dir, 'v1', 'train_feature.csv')
        v1_test = os.path.join(self.data_dir, 'v1', 'test_feature.csv')

        if os.path.exists(v1_train) and os.path.exists(v1_test):
            versions.append('v1')
            print("V1: 基础特征版本")
            print(f"  训练集: {v1_train}")
            print(f"  测试集: {v1_test}")

        # 检查v2版本
        v2_train = os.path.join(self.data_dir, 'v2', 'train_full_v2.csv')
        v2_test = os.path.join(self.data_dir, 'v2', 'test_full_v2.csv')

        if os.path.exists(v2_train) and os.path.exists(v2_test):
            versions.append('v2')
            print("\nV2: 增强特征版本")
            print(f"  训练集: {v2_train}")
            print(f"  测试集: {v2_test}")

        return versions

    def load_v1(self):
        """加载v1版本数据"""
        print("加载V1版本数据...")

        train_path = os.path.join(self.data_dir, 'v1', 'train_feature.csv')
        test_path = os.path.join(self.data_dir, 'v1', 'test_feature.csv')

        if not os.path.exists(train_path):
            raise FileNotFoundError(f"V1训练数据不存在: {train_path}")
        if not os.path.exists(test_path):
            raise FileNotFoundError(f"V1测试数据不存在: {test_path}")

        train_df = pd.read_csv(train_path)
        test_df = pd.read_csv(test_path)

        print(f"  [OK] 训练集: {train_df.shape}")
        print(f"  [OK] 测试集: {test_df.shape}")

        # 加载特征重要性（如果存在）
        feature_importance = None
        importance_path = os.path.join(self.data_dir, 'v1', 'feature_importance.csv')
        if os.path.exists(importance_path):
            try:
                feature_importance = pd.read_csv(importance_path)
                print(f"  [OK] 特征重要性: {feature_importance.shape}")
            except Exception as e:
                print(f"  [WARNING] 特征重要性加载失败: {e}")

        return {
            'version': 'v1',
            'train': train_df,
            'test': test_df,
            'train_shape': train_df.shape,
            'test_shape': test_df.shape,
            'feature_importance': feature_importance
        }

    def load_v2(self):
        """加载v2版本数据"""
        print("加载V2版本数据...")

        train_path = os.path.join(self.data_dir, 'v2', 'train_full_v2.csv')
        test_path = os.path.join(self.data_dir, 'v2', 'test_full_v2.csv')

        if not os.path.exists(train_path):
            raise FileNotFoundError(f"V2训练数据不存在: {train_path}")
        if not os.path.exists(test_path):
            raise FileNotFoundError(f"V2测试数据不存在: {test_path}")

        train_df = pd.read_csv(train_path)
        test_df = pd.read_csv(test_path)

        print(f"  [OK] 训练集: {train_df.shape}")
        print(f"  [OK] 测试集: {test_df.shape}")

        return {
            'version': 'v2',
            'train': train_df,
            'test': test_df,
            'train_shape': train_df.shape,
            'test_shape': test_df.shape
        }

    def load_latest(self):
        """加载最新版本数据（优先v2）"""
        print("加载最新版本数据...")

        try:
            return self.load_v2()
        except FileNotFoundError:
            print("V2版本不存在，尝试加载V1版本...")
            return self.load_v1()

    def extract_features(self, data_dict):
        """从完整数据中提取特征和目标"""
        version = data_dict['version']
        train_df = data_dict['train'].copy()
        test_df = data_dict['test'].copy()

        # 提取ID
        train_ids = train_df['Id'] if 'Id' in train_df.columns else None
        test_ids = test_df['Id'] if 'Id' in test_df.columns else None

        # 分离训练集的特征和目标
        if 'SalePrice' in train_df.columns:
            y_train = train_df['SalePrice']
            X_train = train_df.drop(['SalePrice', 'Id'], axis=1, errors='ignore')
        else:
            y_train = None
            X_train = train_df.drop(['Id'], axis=1, errors='ignore')

        # 提取测试集特征
        X_test = test_df.drop(['Id'], axis=1, errors='ignore')

        result = {
            'version': version,
            'X_train': X_train,
            'X_test': X_test,
            'y_train': y_train,
            'train_ids': train_ids,
            'test_ids': test_ids,
            'feature_count': X_train.shape[1],
            'train_samples': X_train.shape[0],
            'test_samples': X_test.shape[0]
        }

        # 添加v1特有的特征重要性
        if 'feature_importance' in data_dict and data_dict['feature_importance'] is not None:
            result['feature_importance'] = data_dict['feature_importance']

        return result

    def get_data_info(self, version='v2'):
        """获取数据信息"""
        print(f"\n{version.upper()} 版本数据信息:")
        print("="*40)

        if version == 'v1':
            v1_data = self.load_v1()
            train_df, test_df = v1_data['train'], v1_data['test']
        else:
            v2_data = self.load_v2()
            train_df, test_df = v2_data['train'], v2_data['test']

        print(f"训练集:")
        print(f"  形状: {train_df.shape}")
        print(f"  列数: {len(train_df.columns)}")

        print(f"\n测试集:")
        print(f"  形状: {test_df.shape}")
        print(f"  列数: {len(test_df.columns)}")

        # 计算缺失率
        print(f"\n训练集缺失率报告 (前10个有缺失的列):")
        train_missing = train_df.isnull().sum()
        train_missing_pct = (train_missing / len(train_df) * 100).round(2)
        missing_cols = train_missing_pct[train_missing_pct > 0].sort_values(ascending=False).head(10)

        if len(missing_cols) > 0:
            for col, pct in missing_cols.items():
                print(f"  {col}: {pct}%")
        else:
            print("  无缺失值")

        print(f"\n测试集缺失率报告 (前10个有缺失的列):")
        test_missing = test_df.isnull().sum()
        test_missing_pct = (test_missing / len(test_df) * 100).round(2)
        test_missing_cols = test_missing_pct[test_missing_pct > 0].sort_values(ascending=False).head(10)

        if len(test_missing_cols) > 0:
            for col, pct in test_missing_cols.items():
                print(f"  {col}: {pct}%")
        else:
            print("  无缺失值")

        # 特征数（排除ID和SalePrice）
        if 'Id' in train_df.columns:
            train_features = train_df.drop(['Id'], axis=1)
        else:
            train_features = train_df

        if 'SalePrice' in train_features.columns:
            train_features = train_features.drop(['SalePrice'], axis=1)

        print(f"\n特征数量: {len(train_features.columns)}")


def main():
    """主函数 - 简化的交互式访问"""
    print("="*60)
    print("简化版数据访问工具")
    print("="*60)

    manager = SimpleDataManager()

    while True:
        print("\n请选择操作:")
        print("1. 列出可用版本")
        print("2. 加载V1版本")
        print("3. 加载V2版本")
        print("4. 加载最新版本")
        print("5. 查看数据信息")
        print("0. 退出")

        try:
            choice = input("\n请选择 (0-5): ").strip()

            if choice == "0":
                print("退出数据访问工具")
                break
            elif choice == "1":
                manager.list_versions()
            elif choice == "2":
                v1_data = manager.load_v1()
                train_df, test_df = v1_data['train'], v1_data['test']
                print(f"V1数据加载完成")
            elif choice == "3":
                v2_data = manager.load_v2()
                train_df, test_df = v2_data['train'], v2_data['test']
                print(f"V2数据加载完成")
            elif choice == "4":
                latest_data = manager.load_latest()
                train_df, test_df = latest_data['train'], latest_data['test']
                print(f"最新数据加载完成")
            elif choice == "5":
                version = input("请输入版本 (v1/v2, 默认v2): ").strip() or 'v2'
                manager.get_data_info(version)
            else:
                print("无效选择，请重新输入")

        except KeyboardInterrupt:
            print("\n\n用户中断操作")
            break
        except Exception as e:
            print(f"\n操作失败: {e}")


# 使用示例
if __name__ == "__main__":
    # 快速使用示例
    print("使用示例:")
    print("-" * 30)

    # 初始化管理器
    manager = SimpleDataManager()

    # 列出版本
    versions = manager.list_versions()

    # 加载V2数据并展示部分数据
    if 'v2' in versions:
        print("\n加载V2数据示例:")
        v2_data = manager.load_v2()
        train_df = v2_data['train']
        test_df = v2_data['test']

        print(f"\n训练集前3行数据预览:")
        print(train_df.head(3))

        print(f"\n测试集前3行数据预览:")
        print(test_df.head(3))

        print(f"\n数据概览:")
        print(f"训练集形状: {train_df.shape}")
        print(f"测试集形状: {test_df.shape}")

    # 获取数据信息和缺失率报告
    print("\n" + "="*60)
    main()